# import module
import TextGenerator
import torch  # 导入 PyTorch 库，用于深度学习任务。
import gradio as gr  # 导入 Gradio 库，用于构建交互式界面。
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.generation import GenerationConfig



# # 定义模型ID
# model_id = "shenzhi-wang/Llama3-8B-Chinese-Chat"
# # model_id = "hf-models/Llama-2-7b-chat-hf"

# # 加载分词器和模型
# tokenizer = AutoTokenizer.from_pretrained(model_id)
# model = AutoModelForCausalLM.from_pretrained(
#     model_id, torch_dtype="auto", device_map="auto"
# )

# model.eval()

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-7B-Chat", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", device_map="auto", trust_remote_code=True).eval()

# 创建类实例
text_generator = TextGenerator.TextGenerator(model, tokenizer)


# 创建 Gradio 接口
demo = gr.Blocks()  # 在全局范围内定义 demo22

with demo:
    with gr.Tab("续写"):
        with gr.Column():
            input_text_zhengwen = gr.Textbox(lines=2, placeholder="请输入要续写的文本", label="正文")
            input_text_dagang = gr.Textbox(lines=2, placeholder="请输入续写的主题或方向", label="大纲")
            output_textxuxie = gr.Textbox(lines=4, placeholder="生成的续写文本", label="续写结果")
            btn_xuxie = gr.Button("生成")
            btn_xuxie.click(text_generator.continue_writting, inputs=[input_text_zhengwen, input_text_dagang], outputs=output_textxuxie)

    with gr.Tab("头脑风暴"):
        with gr.Column():
            prompt_input = gr.Textbox(lines=2, placeholder="请输入主题", label="主题")
            output_texttounao = gr.Textbox(lines=4, placeholder="生成的头脑风暴文本", label="生成结果")
            btn_tounao = gr.Button("生成")
            btn_tounao.click(text_generator.generate_text_tounao_baofeng, inputs=[prompt_input], outputs=output_texttounao)
    
    with gr.Tab("找错别字"):
        with gr.Column():
            prompt_content = gr.Textbox(lines=4, placeholder="请输入正文", label="正文")
            output_text_bug = gr.Textbox(lines=2, placeholder="", label="生成结果")
            btn_findbug = gr.Button("查找错别字")
            btn_findbug.click(text_generator.generate_text_find_error_bug, inputs=[prompt_content], outputs=output_text_bug)

    with gr.Tab("转换风格"):
        with gr.Column():
            # examples = [["小白风格"], ["玄幻风格"], ["都市风格"]]
            prompt_content = gr.Textbox(lines=4, placeholder="请输入正文", label="正文")
            change_type = gr.Dropdown(
                choices=["小白风格", "玄幻风格", "都市风格"],
                label="选择转换风格"
            )
            output_text_bug = gr.Textbox(lines=2, placeholder="", label="生成结果")
            btn_change_type = gr.Button("转换风格")
            btn_change_type.click(text_generator.generate_text_change, inputs=[prompt_content,change_type], outputs=output_text_bug)

# 启动界面
if __name__ == "__main__":
   demo.queue().launch()